# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tensorflow as tf
from tensorflow.python.tools import freeze_graph
from model import select_model
import numpy as np

tf.app.flags.DEFINE_string('model_path','./LMTCNN-1-1-run-7133/checkpoint-49999',
	'valiation and testing set directory')
FLAGS = tf.app.flags.FLAGS
# ckpt_path  = '/LMTCNN-1-1-run-7133/checkpoint-1000'

def main(): 
    model_fn = select_model('LMTCNN-1-1')
    tf.reset_default_graph()
    # 定义网络的输入节点
    image_holder = tf.placeholder(tf.float32, shape=[None, 227,227,3], name="image")
    agelogits, genderlogits = model_fn(8, image_holder,2, image_holder, 1, False)
    predict_gender = tf.argmax(genderlogits, axis=1, output_type=tf.int32, name="gender")
    # 定义网络的输出节点
    
    with tf.Session() as sess:
        #保存图，在./pb_model文件夹中生成model.pb文件
        # model.pb文件将作为input_graph给到接下来的freeze_graph函数
        tf.train.write_graph(sess.graph_def, './pb_model', 'model.pb')    # 通过write_graph生成模型文件
        freeze_graph.freeze_graph(
		        input_graph='./pb_model/model.pb',   # 传入write_graph生成的模型文件
		        input_saver='',
		        input_binary=False, 
		        input_checkpoint=FLAGS.model_path,  # 传入训练生成的checkpoint文件
		        output_node_names='ageoutput/ageoutput,gender',  # 与定义的推理网络输出节点保持一致
		        restore_op_name='save/restore_all',
		        filename_tensor_name='save/Const:0',
		        output_graph='./pb_model/lmtcnn_1.pb',   # 改为需要生成的推理网络的名称
		        clear_devices=False,
		        initializer_nodes='')
    print("done")

if __name__ == '__main__': 
    main()